import math

POSINF=1e300                            # +ve infinity
NEGINF=-1e300                           # -ve infinity
POSZERO=1e-300                          # +0
NEGZERO=-1e-300                         # -0

class DefaultDict(dict):
  """
  Creates a dictionary with a default value function.
  Normally a None value returned by the default value function causes a
  key not found exception. However, if the none_valid argument is set to true
  then this feature is over-ridden, i.e. None values are returned by the
  dictionary.
  """
  def __init__(self, *args, **kwargs):
    """
    default - a function which takes the idx entry as the argument
    none_valid - defaults to False which indicates that if the default fn
                 returns None then a key error should be raised
    """
    # extract the default value from the keywrord arguments
    if "default" not in kwargs:
      raise "DefaultDict needs a default keyword argument"
    self.default = kwargs["default"]
    del kwargs["default"]
    
    # check for the none_valid argument
    if "none_valid" in kwargs:
      self.none_valid = True
      del kwargs["none_valid"]
    else:
      self.none_valid = False

    # pass the rest on to the dictionary base-type
    dict.__init__(self, *args, **kwargs)

  def copy(self):
    newdict = DefaultDict(default=self.default, none_valid=self.none_valid)
    for k,v in self.iteritems():
      newdict[k] = v.copy()
    return newdict
  
  def __getitem__(self, idx):
    try:
      return dict.__getitem__(self, idx)
    except KeyError:
      # if there is no default function then raise the exception
      if self.default is None:
        raise
      # try the default function
      val = self.default(idx)
      # re-raise the error if the default function returns None and
      # None is not a valid value
      if (val is None) and (not self.none_valid):
        raise
      # otherwise store the default in the index location and return
      # that value
      self[idx] = val
      return val
    
  def nodefault(self):
    """Disable the default function."""
    self.default = None

# loosely inspired by a Counter in Dan Klein's CS 188 (UC Berkeley) course
# that I encountered in Fall 2006.
class Counter(dict):
  """
  If the keyword default is passed as an argument a default function
  can be used to generate a value on key not found
  """
  def __init__(self, *args, **kwargs):
    # extract the default value from the keywrord arguments
    if "default" in kwargs:
      self.default = kwargs["default"]
      del kwargs["default"]
    else:
      self.default = 0

    # pass the rest on to the dictionary base-type
    dict.__init__(self, *args, **kwargs)

  def nodefault(self):
    self.default = None

  def newdefault(self, val):
    self.default = val
  
  def __getitem__(self, idx):
    try:
      return dict.__getitem__(self, idx)
    except KeyError:
      if self.default is not None:
        return self.default
      else:
        raise

  def __abs__(self):
    """
    The L_2 norm
    """
    tot = 0
    for event, cnt in self.iteritems():
      tot += cnt*cnt
    return math.sqrt(tot)
  
  def __add__(self, arg):
    ret = Counter(self.iteritems())
    if type(arg) is Counter:
      for e,c in arg.iteritems():
        ret[e] += c
    else:
      for e in ret:
        ret[e] += arg
    return ret

  def __iadd__(self, arg):
    if type(arg) is Counter:
      for e,c in arg.iteritems():
        self[e] += c
    else:
      for e in self:
        self[e] += arg
    return self

  def __mul__(self, arg):
    """
    Multiply each count by a scalar
    """
    ret = Counter(self.iteritems())
    for e in ret:
      ret[e] *= arg
    return ret

  def __imul__(self, arg):
    """
    Scale by a scalar
    """
    for e in self:
      self[e] *= arg

    return self

  def __div__(self, arg):
    """
    Multiply each count by a scalar
    """
    ret = Counter(self.iteritems())
    for e in ret:
      ret[e] /= arg
    return ret

  def __idiv__(self, arg):
    """
    Scale by a scalar
    """
    for e in self:
      self[e] /= arg

    return self

  def log(self):
    """
    Take a logarithm of all the counts.
    """
    for e in self:
      self[e] = math.log(self[e])

  def exp(self):
    """
    Take a exponential of all the counts.
    """
    for e in self:
      self[e] = math.exp(self[e])
    
  def copy(self):
    """
    Return a copy of this counter with all the values copied as well
    """
    newcnt = Counter(default=self.default)
    for k,v in self.iteritems():
      newcnt[k] = v.copy()
    return newcnt

  
  def total(self):
    tot = 0
    for event, cnt in self.iteritems():
      tot += cnt
    return tot
  
  def normalize(self):
    """
    Divide all the counts of events by the total counts
    """
    tot = self.total()
    if tot == 0.0: return
    for event, cnt in self.iteritems():
      self[event] = float(cnt) / tot

  def max(self):
    bestcnt = None
    for event, cnt in self.iteritems():
      if cnt > bestcnt: bestcnt = cnt
    return bestcnt

  def argmax(self):
    bestevt = bestcnt = None
    for event, cnt in self.iteritems():
      if cnt > bestcnt: bestevt, bestcnt = event, cnt
    return bestevt

  mode = argmax
  
  def min(self):
    worstcnt = None
    for event, cnt in self.iteritems():
      if -cnt > worstcnt: worstcnt = -cnt
    if worstcnt is None:
      return None
    return -worstcnt

  def argmin(self):
    worstevt = worstcnt = None
    for event, cnt in self.iteritems():
      if -cnt > worstcnt: worstevt, worstcnt = event, -cnt
    return worstevt

  def smooth_abs_discount(self, D1=0.5, D2=None):
    """
    Discount the events which occur by subtract an amount from their counts,
    and putting the remaining mass into the "None" event.
    If only D1 is given, then all events counts are subtracted by D1. If D2 is
    also given then events with a count of 2 or higher are subtracted by D2.
    """
    assert((D1 > 0) and (D1 < 1))
    if D2 is None: D2 = D1
    assert((D2 > 0) and (D2 < 2))
    # now subtract the amount from all events and compute leftover
    leftover = 0.0
    for event, count in self.iteritems():
      assert((event is not None) and (count > 0) and (type(count) is int))
      if count == 1:
        self[event] -= D1
        leftover += D1
      else:
        self[event] -= D2
        leftover += D2
    # put the left over mass in the unknown event
    self[None] = leftover

  def prune(self, count):
    """
    Delete all but the top 'count' entries
    """
    if len(self) <= count:
      return
    k = self.keys()
    k.sort(cmp=lambda x,y: int(100*(self[x] - self[y])))
    for i in range(len(self) - count):
      del self[k[i]]

  def trim(self, threshold = 1e-300):
    """
    Delete all keys which have a near 0 count i.e. abs(c) <= threshold
    """
    selected = []
    for k, c in self.iteritems():
      if abs(c) <= threshold:
        selected.append(k)
    for k in selected:
      del self[k]

  def sorted_key_vals(self, reverse=False):
    """
    Returns a list of key, value pairs sorted by value
    """
    key_vals = [x for x in self.iteritems()]
    key_vals.sort(cmp=lambda x,y: 2*int(x[1] > y[1]) - 1, reverse=reverse)
    return key_vals

  def entropy(self):
    """
    Assuming the counter is normalized returns the entropy
    of this distribution, i.e. the number of bits of information
    gained by knowing the correct key
    """
    ent = 0.0
    for k,v in self.iteritems():
      if v == 0:
        continue
      ent += - v * math.log(v, 2)
    assert(ent >= 0)
    return ent

  def kl_divergence(self, cnt2, discount=1e-6):
    """
    Compute the kl divergence between two normalized probability
    distributions.

    discount - amount of probability mass to spread among elements of cnt2
               so that cnt2 always has non-zero probability for each element
               in cnt1
    """
    epsilon = float(discount) / len(self)
    div = 0.0
    for k,v in self.iteritems():
      if v == 0:
        continue
      div += v * math.log( v / (((1 - discount)*cnt2[k]) + epsilon))
    return div

  def median(self):
    """
    Sorts the keys and then returns the smallest key such that the cumulative
    probability is greater than or equal to .5
    """
    keys = self.keys()
    keys.sort()
    cum = 0.0
    for k in keys:
      cum += self[k]
      if cum >= .5:
        return k
    raise ValueError("median called on Counter whose total prob < .5")
  
  def wt_sum(self):
    """
    Weighted sum of the elements of the counter (sum k*v)
    """
    return sum(k*v for k,v in self.iteritems())

  mean = wt_sum
  avg = wt_sum

  def std_dev(self):
    """
    Assuming that this is a probability distribution on numbers. Compute
    the standard deviation
    """
    # sqrt( E (x^2) - E(x) ^ 2)
    return math.sqrt(sum(v*k**2 for k,v in self.iteritems()) \
                     - sum(v*k for k,v in self.iteritems()) ** 2)
  
  def sample(self, rand):
    """
    Assuming that this is a probability distribution (with total 1).
    This function will randomly sample an element.
    """
    choice = rand.random()
    s = 0.0
    for k,v in self.iteritems():
      s += v
      if s > choice:
        break
    else:
      import pdb
      pdb.set_trace()
      assert(False)                     # sum doesn't add to one
    return k
  
  def logadd(self, item, logval):
    """
    self[item] =  log(exp(self[item]) + exp(logval))
    """
    x, y = self[item], logval
    x, y = max(x,y), min(x,y)
    self[item] = x + math.log(1+math.exp(y-x))

  def lognormalize(self):
    """
    Assuming that we are storing the log of the weights. This converts the
    counter into a probability distribution using the weights.
    """
    m = max(self.itervalues())
    tot = sum(math.exp(v-m) for v in self.itervalues())
    for k, v in self.iteritems():
      self[k] = float(math.exp(v-m))/tot
    
# to dynamically a load a class from a module  
def importfrom(modname, classname=None):
  # extract the module name and class name
  if (classname is None) and (modname.rfind(".") >= 0):
    cls_pos = modname.rfind(".")
    classname = modname[cls_pos+1:]
    modname = modname[:cls_pos]
  module = __import__(modname, None, None, classname)
  classobj = getattr(module, classname)
  return classobj

def parse_fncall(fncall):
  """
  Given a function call string like punini.chunker.Dtree?N=4,D=5,M=.8  returns
  (punini.chunker.Dtree, {N:int(4), D:int(5), M:float(.8)})
  """
  qpos = fncall.find("?")
  # are there any parameters
  if qpos == -1:
    return (fncall, {})
  
  fnname = fncall[:qpos]
  fnpara = {}
  for par in fncall[qpos+1:].split(","):
    lhs, rhs = par.split("=")
    # by default, try the parameter as a string
    if rhs.lower() in ["true", "false"]:
      rhs_val = bool(rhs)
    else:
      rhs_val = str(rhs)
    # is this floating point?
    try:
      rhs_val = float(rhs)
    except ValueError:
      pass
    # is this integer?
    try:
      rhs_val = int(rhs)
    except ValueError:
      pass
    fnpara[lhs] = rhs_val

  return (fnname, fnpara)

def log_add_exp(x,y):
  """
  returns log(exp(x) + exp(y))
  """
  t = math.exp(x) + math.exp(y)
  if t == 0.0:
    return max(x, y)
  else:
    return math.log(t)

def memoize(memo_table):
  """
  Usage:
  -----
  
  cache = {}
  @memoize(cache)
  def foo(..):
    ...

  This will store the memoized calls to function "foo" in hash table "cache"

  """
  def decorator(func):
    def wrapper(*args, **kwargs):
      fnkey = (func, tuple(args), tuple(kwargs.iteritems()))
      try:
        return memo_table[fnkey]
      except KeyError:
        val = func(*args, **kwargs)
        memo_table[fnkey] = val
        return val
    # safe decorator practice -- copy name, attribute and doc string from
    # original function
    wrapper.__name__ = func.__name__
    wrapper.__dict__ = func.__dict__
    wrapper.__doc__ = func.__doc__
    return wrapper
  return decorator

def avgresults(results):
  """
  Given a list of r.v. -> val maps this returns the average and standard
  deviation of all the random variables
  """
  # compute the moments of all the keys in the results
  moments = DefaultDict(default=lambda x: [0,0,0])
  for result in results:
    for k,v in result.iteritems():
      m = moments[k]
      m[0] += 1
      m[1] += v
      m[2] += v*v
  # now compute the average and std. dev. of these keys
  retlist = []
  for k,m in moments.iteritems():
    av = float(m[1])/m[0]
    var = (float(m[2])/m[0]) - (av * av)
    sd = math.sqrt(var)
    retlist.append((k, av, sd))
  return retlist
      
class CircularList:
  """
  Allows items to be appended which overwrites the initial item after the
  buffer is full.
  """
  def __init__(self, maxsize):
    self.maxsize = maxsize
    self.pos = 0
    self.list = []
  def append(self, elem):
    if len(self.list) < self.maxsize:
      self.list.append(elem)
    else:
      self.list[self.pos] = elem
      self.pos = (self.pos + 1) % self.maxsize
  def __iter__(self):
    return self.list.__iter__()
  def __str__(self):
    return self.list.__str__()
    
